package com.tianji.aigc.ChatMemory;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.collection.ListUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.aigc.entity.RedisMessage;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import javax.swing.plaf.ListUI;
import java.util.List;
import java.util.stream.Stream;

/**
 * Redis会话记忆处理类
 */
@Component
public class RedisChatMemory implements ChatMemory {
    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 保存会话对应的消息到Redis中
     * @param conversationId
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        // 1.消息为空,则返回
        if(CollUtil.isEmpty(messages)){
            return;
        }
        // 2.否则保存消息
        // 2.1 获取Redis Key
        String key = getRedisKey(conversationId);
        // 2.2将Message消息转为JSON字符串
        List<String> messageStrList = messages.stream().map(message -> {
            //拷贝Message对象为自定义的RedisMessage对象
            RedisMessage redisMessage = BeanUtil.toBean(message, RedisMessage.class);
            //设置消息内容字符串，否则为空
            redisMessage.setTextContent(message.getText());
            //设置Tool相关属性的值
            if(message instanceof AssistantMessage assistantMessage) {
                redisMessage.setToolCalls(assistantMessage.getToolCalls());
            } else if(message instanceof ToolResponseMessage toolResponseMessage) {
                redisMessage.setToolResponses(toolResponseMessage.getResponses());
            }
            return JSONUtil.toJsonStr(redisMessage);
        }).toList();
        // 2.3将消息JSON字符串列表添加到Redis的List中
        redisTemplate.opsForList().rightPushAll(key, messageStrList);
    }

    /**
     * 从Redis中查询某会话下的指定条数的消息列表
     * @param conversationId
     * @param lastN
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 1.参数校验
        if(lastN<=0){
            return List.of();
        }
        // 2.获取redisKey，组成结构：CHAT:用户id:sessionId
        String key = this.getRedisKey(conversationId);
        // 3.从Redis的List中获取消息字符串列表
        List<String> messageStrList = redisTemplate.opsForList().range(key, 0, lastN);
        if(CollUtil.isEmpty(messageStrList)) {
            return List.of();
        }
        // 4.将消息JSON字符串转为Message
        List<Message> messageList = messageStrList.stream().map(x -> {
            //拷贝Message对象为自定义的RedisMessage对象
            RedisMessage redisMessage = JSONUtil.toBean(x, RedisMessage.class);
            Message message = null;
            if(redisMessage.getMessageType().equals(MessageType.ASSISTANT)) {
                message = new AssistantMessage(redisMessage.getTextContent(),redisMessage.getMetadata(),redisMessage.getToolCalls(),redisMessage.getMedia());
            } else if(redisMessage.getMessageType().equals(MessageType.TOOL)) {
                message = new ToolResponseMessage(redisMessage.getToolResponses(),redisMessage.getMetadata());
            } else if(redisMessage.getMessageType().equals(MessageType.USER)) {
                message = new UserMessage(redisMessage.getMessageType(), redisMessage.getTextContent(),redisMessage.getMedia(),redisMessage.getMetadata());
            } else {
                message = new SystemMessage(redisMessage.getTextContent());
            }
            return message;
        }).toList();
        return messageList;
    }
    /**
     * 删除Redis中某会话下的所有消息
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        String key = this.getRedisKey(conversationId);
        redisTemplate.delete(key);
    }

    /**
     * 拼接RedisKey
     * @param conversationId
     * @return
     */
    private String getRedisKey(String conversationId){
        return "CHAT:" + conversationId;
    }
}
